import trimesh
import math
import scipy
import os
import numpy as np
import json
import pickle
from copy import deepcopy
from perception import DepthImage, CameraIntrinsics
from autolab_core import RigidTransform
import multiprocessing as mp
# from save_dict import save_new_npy


def change_ray_color(ray, color):
    colors = np.ones((len(ray.entities), 3))
    colors_1 = (colors * color).astype(np.uint8)  # approach
    ray.colors = colors_1


def show(scene):
    scene.show()


def save_dict(dict_, mask_data,k):
    try:
        np.savez('./mask_label_grasp/image_{:05d}.npz'.format(k), dict_)
        np.save('./mask_label_grasp/mask_{:05d}.npy'.format(k), mask_data)
        print('save finished')
    except:
        print('save failed')
        pass


def save_arr(arr, k):
    try:
        arr = arr.astype(np.float16)
        np.save('./mask_label_grasp/image_{:05d}.npy'.format(k), arr)
        # save_new_npy(k)
        # print('save finished')
    except:
        print('save failed')
        pass


def generate_label(k, suction_gripper, base, finger, bin_mesh, obj_ids, obj_poses, image_data, camera_intr_,
                   camera_pose, modal_mask):
    print('k is ', k)
    if os.path.exists('./mask_label_grasp/image_{:05d}.npy'.format(k)):
        print('./mask_label_grasp/image_{:05d}.npy already exists'.format(k))
        return
    scene = trimesh.Scene()
    scene.add_geometry(bin_mesh)

    depth_image = DepthImage(image_data, frame='camera')
    camera_intr = CameraIntrinsics(frame='camera', fx=camera_intr_[0], fy=camera_intr_[1], cx=camera_intr_[2],
                                   cy=camera_intr_[3], skew=camera_intr_[4], height=camera_intr_[5],
                                   width=camera_intr_[6])
    point_normal_cloud = depth_image.point_normal_cloud(camera_intr)
    point_data = point_normal_cloud.points.data
    T_matrix = trimesh.transformations.translation_matrix(camera_pose[:3])
    rotation_matrix = trimesh.transformations.quaternion_matrix(camera_pose[3:])
    matrix = trimesh.transformations.concatenate_matrices(T_matrix, rotation_matrix)
    matrix = np.asanyarray(matrix, order='C', dtype=np.float64)
    point_data = trimesh.transformations.transform_points(point_data.transpose(1, 0), matrix)
    foreground_mask = point_data[:, 2] > 0.01
    point_data = point_data[foreground_mask]

    point_num = point_data.shape[0]
    if parallel_jaw_grasp:
        point_label_arr = np.zeros([point_num, 24])
    elif suction:
        point_label_arr = np.zeros([point_num, 6])
    else:
        assert False
    # print(point_label_arr.shape)


    # reshape modal_mask
    modal_mask = modal_mask.reshape(-1, 16).T
    modal_mask = modal_mask[:, foreground_mask]
    max_objs_per_scene = modal_mask.shape[0]

    point_grasp_dict = {}

    point_tree = scipy.spatial.KDTree(point_data)
    point_cloud = trimesh.PointCloud(point_data, colors=[0, 255, 0])
    # point_cloud.show()

    obj_ids_per_scene, obj_poses_per_scene = obj_ids[(k % 200) // 2], obj_poses[(k % 200) // 2]
    obj_dict = {}
    mesh_dict = {}

    for obj_id, obj_pose, idx in zip(obj_ids_per_scene, obj_poses_per_scene, range(max_objs_per_scene)):
        if obj_id == 4294967295:
            continue
        # print(obj_id, obj_pose)
        obj_key = metadata['obj_ids'][str(obj_id[0])]
        obj_path = metadata['meshes'][obj_key]
        obj_path = os.path.join('../grasp_suction_dataset/object_models/meshes/', obj_path.split('/')[-1])
        obj = trimesh.load_mesh(obj_path)
        trimesh.smoothing.filter_humphrey(obj)
        T_matrix = trimesh.transformations.translation_matrix(obj_pose[:3])
        rotation_matrix = trimesh.transformations.quaternion_matrix(obj_pose[3:])
        matrix = trimesh.transformations.concatenate_matrices(T_matrix, rotation_matrix)
        obj.apply_transform(matrix)
        scene.add_geometry(obj)
        # for debug
        obj_point_index = np.where(modal_mask[idx] > 0)[0]
        obj_dict[obj_key] = (matrix, obj_point_index)
        mesh_dict[obj_key] = obj

    collision_manager, _ = trimesh.collision.scene_to_collision(scene)
    signed_distance = trimesh.proximity.signed_distance(bin_mesh, point_data)
    mask = abs(signed_distance) < 0.0005
    points_query_bin = [i for i in range(len(mask)) if mask[i]]
    # pc = trimesh.PointCloud(point_data[points_query_bin], colors=[0, 0, 255])
    # pc.show()
    for index in points_query_bin:
        if index not in point_grasp_dict.keys():
            if parallel_jaw_grasp:
                point_grasp_dict[index] = (0, 0.0, np.identity(4).astype(np.float16), 0.0, 0.0, 0, 0.0, 0.0, 0.0) # label[-3] is distance and label[-4] is objectness label
                point_label_arr[index] = [0, 0.0, *np.identity(4).astype(np.float16).reshape(16), 0.0, 0.0, 0, 0.0,0.0,0.0]#label[-1] new_width by offset label[-2] center offset_distance along grasp axis
            elif suction:
                point_grasp_dict[index] = (0, 0.0, -1, -1, -1, 0)
                point_label_arr[index] = [0, 0.0, -1, -1, -1, 0]

    point_center_direction = np.zeros((modal_mask.shape[1], 3))
    # scene_tmp = trimesh.Scene()
    # pointcloud = trimesh.PointCloud(point_data, colors=[0, 255, 0])
    # scene_tmp.add_geometry(pointcloud)
    for key, (value, obj_point_index) in obj_dict.items():
        key_ori = key
        key = key.split('~')[1]
        point_center_direction[obj_point_index] = value[:3, 3] - point_data[obj_point_index]

        # print('point_direction :', point_center_direction[obj_point_index])
        # ray_visualize = trimesh.load_path(
        #     np.hstack((point_data[obj_point_index], point_data[obj_point_index] + point_center_direction[obj_point_index])).reshape(-1, 2, 3))
        # scene_tmp.add_geometry(ray_visualize)
        # continue
        # load the badpoints for negtive points
        badpoints = np.load(grasp_path + '{}/badpoints.npy'.format(key))
        if badpoints.shape[0] > 0:
            badpoints = trimesh.transformations.transform_points(badpoints, value)
            points_query = point_tree.query_ball_point(badpoints, r=0.0025, p=2)  # 0.005 for grasp
            points_query = [item for sublist in points_query for item in sublist]
            points_query = list(set(points_query))
            # the default setting let the points_query not appear in background or other object
            if len(points_query) > 0:
                signed_distance = trimesh.proximity.signed_distance(mesh_dict[key_ori], point_data[points_query])
                mask = abs(signed_distance) < 0.0005
                points_query = [points_query[i] for i in range(len(mask)) if mask[i]]
                for index in points_query:
                    if index not in point_grasp_dict.keys() or 0 > point_grasp_dict[index][0]:
                        if parallel_jaw_grasp:
                            point_grasp_dict[index] = (0, 0.0, np.identity(4).astype(np.float16), 0.0, 0.0, 1, 0.0, 0.0, 0.0)
                            point_label_arr[index] = [0, 0.0, *np.identity(4).astype(np.float16).reshape(16), 0.0, 0.0,
                                                      1, 0.0, 0.0, 0.0]
                        elif suction:
                            point_grasp_dict[index] = (0, 0.0, -1, -1, -1, 1)
                            point_label_arr[index] = [0, 0.0, -1, -1, -1, 1]
        files = os.listdir(grasp_path + '{}'.format(key))
        for file_name in files:
            if os.path.splitext(file_name)[0].startswith('final'):  # final for grasp
                with open(grasp_path + '{}/{}'.format(key, file_name), 'rb') as f:
                    grasps = pickle.load(f)
                    #                    pruned_grasps = GraspSampler.nms_grasps(grasps,threshold=0.01)
                    #                    for grasp in pruned_grasps:
                    if parallel_jaw_grasp:
                        for grasp in grasps:
                            is_collision_flag = False
                            pose = RigidTransform(grasp.approached_full_axis, grasp.center, from_frame='gripper',
                                                  to_frame='obj')
                            transform = trimesh.transformations.concatenate_matrices(value, pose.matrix)

                            z_vec = transform[:3, 0]   #z_vec = transform[:3, 0]
                            # print(z_vec)
                            if z_vec[2] > -math.sin(math.pi / 12):
                                # print('test a single grasp angle')
                                is_collision_flag = True
                                # continue
                            base_copy = deepcopy(base)
                            base_copy.apply_transform(transform)
                            finger_left_copy = deepcopy(finger)
                            T_finger = trimesh.transformations.translation_matrix(
                                [-(grasp.real_jaw_width_ / 2 + 0.002), 0, 0])
                            R_finger_left = trimesh.transformations.euler_matrix(np.pi / 2, np.pi / 2, 0, 'rxyz')
                            Trans_finger_left = trimesh.transformations.concatenate_matrices(transform,
                                                                                             R_finger_left,
                                                                                             T_finger)

                            finger_left_copy.apply_transform(Trans_finger_left)
                            finger_right_copy = deepcopy(finger)
                            R_finger_right = trimesh.transformations.euler_matrix(np.pi / 2, np.pi / 2, np.pi,
                                                                                  'rxyz')
                            Trans_finger_right = trimesh.transformations.concatenate_matrices(transform,
                                                                                              R_finger_right,
                                                                                              T_finger)
                            finger_right_copy.apply_transform(Trans_finger_right)

                            if is_collision_flag:
                                is_collision = True
                            else:
                                is_collision = collision_manager.in_collision_single(finger_left_copy,
                                                                                     return_names=False) \
                                               or collision_manager.in_collision_single(finger_right_copy,
                                                                                        return_names=False) \
                                               or collision_manager.in_collision_single(base_copy,
                                                                                        return_names=False)

                            if not is_collision:
                                grasp_point = grasp.grasp_point
                                grasp_point = trimesh.transformations.transform_points(grasp_point.reshape(1,3), value)
                                # print('distance is :', np.linalg.norm(transform[:3, 3] - grasp_point))
                                # v1 = transform[:3, 3] - grasp_point
                                # v2 = z_vec
                                # angle = v1.dot(v2)/(np.linalg.norm(v1)*np.linalg.norm(v2))
                                # print(angle)

                                # exit()
                                # contact_points_cloud = trimesh.PointCloud(contact_points, colors=[255, 0, 0])
                                points_query = point_tree.query_ball_point(grasp_point, r=0.0025, p=2)[0]
                                # if len(points_query[0]) > 2:
                                #     if len(points_query[1]) > 2:
                                #         points_query = points_query[0] + points_query[1]
                                #     else:
                                #         points_query = points_query[0]
                                # elif len(points_query[1]) > 2:
                                #     points_query = points_query[1]
                                # else:
                                #     points_query = []

                                if len(points_query) > 0:
                                    # the default setting let the points_query not appear in background or other
                                    # object
                                    signed_distance = trimesh.proximity.signed_distance(mesh_dict[key_ori],
                                                                                        point_data[points_query])
                                    mask = abs(signed_distance) < 0.0005
                                    points_query = [points_query[i] for i in range(len(mask)) if mask[i]]
                                    for index in points_query:
                                        if index not in point_grasp_dict.keys() or grasp.grasp_score.sum() > \
                                                point_grasp_dict[index][1] or point_grasp_dict[index][0] == 0:

                                            print('extend pad ****************', grasp.real_jaw_width_ - grasp.jaw_width,
                                                  'extend pad ****************')
                                            extend_pad = grasp.real_jaw_width_ - grasp.jaw_width
                                            grasp_width_new_ = grasp.grasp_width_new + extend_pad
                                            print('new_width',grasp_width_new_)
                                            print('ori width',grasp.real_jaw_width_)

                                            point_grasp_dict[index] = (
                                                1, grasp.grasp_score.sum(), transform.astype(np.float16), grasp.jaw_width_,
                                                grasp.real_jaw_width_, 1, grasp.approach_distance, grasp.offset_distance, grasp_width_new_)
                                            point_label_arr[index] = [
                                                1, grasp.grasp_score.sum(), *transform.astype(np.float16).reshape(16), grasp.jaw_width_,
                                                grasp.real_jaw_width_, 1, grasp.approach_distance, grasp.offset_distance, grasp_width_new_]
                            else:
                                grasp_point = grasp.grasp_point
                                # print(grasp_point)
                                grasp_point = trimesh.transformations.transform_points(grasp_point.reshape(1,3), value)
                                # grasp_point_cloud = trimesh.PointCloud(grasp_point, colors=[255, 0, 0])
                                # scene.add_geometry(grasp_point_cloud)
                                # # scene_tmp = trimesh.Scene([grasp_point_cloud, point_cloud])
                                # scene.show()
                                # exit()
                                # contact_points_cloud = trimesh.PointCloud(contact_points, colors=[255, 0, 0])
                                points_query = point_tree.query_ball_point(grasp_point, r=0.0025, p=2)[0]
                                # if len(points_query[0]) > 3:
                                #     if len(points_query[1]) > 3:
                                #         points_query = points_query[0] + points_query[1]
                                #     else:
                                #         points_query = points_query[0]
                                # elif len(points_query[1]) > 3:
                                #     points_query = points_query[1]
                                # else:
                                #     points_query = []
                                if len(points_query) > 0:
                                    # the default setting let the points_query not appear in background or other
                                    # object
                                    signed_distance = trimesh.proximity.signed_distance(mesh_dict[key_ori],
                                                                                        point_data[points_query])
                                    mask = abs(signed_distance) < 0.0005
                                    points_query = [points_query[i] for i in range(len(mask)) if mask[i]]
                                    for index in points_query:
                                        if index not in point_grasp_dict.keys() or (
                                                grasp.grasp_score.sum() > point_grasp_dict[index][1] and
                                                point_grasp_dict[index][0] == 0):
                                            print('extend pad ****************', grasp.real_jaw_width_ - grasp.jaw_width,
                                                  'extend pad ****************')
                                            extend_pad = grasp.real_jaw_width_ - grasp.jaw_width
                                            grasp_width_new_ = grasp.grasp_width_new + extend_pad
                                            print('new_width', grasp_width_new_)
                                            print('ori width', grasp.real_jaw_width_)
                                            point_grasp_dict[index] = (
                                                0, grasp.grasp_score.sum(), transform.astype(np.float16), grasp.jaw_width_,
                                                grasp.real_jaw_width_, 1, grasp.approach_distance, grasp.offset_distance, grasp_width_new_)
                                            point_label_arr[index] = [
                                                0, grasp.grasp_score.sum(), *transform.astype(np.float16).reshape(16), grasp.jaw_width_,
                                                grasp.real_jaw_width_, 1, grasp.approach_distance, grasp.offset_distance, grasp_width_new_]

                            if debug_vis and not is_collision:
                                # if debug_vis:
                                scene_ = deepcopy(scene)
                                ray_origins = grasp_point#np.array([0., 0., 0.])
                                pc_center = trimesh.PointCloud(transform[:3, 3].reshape(-1, 3), colors=[0, 0, 255])
                                ray_directions = z_vec
                                ray_visualize = trimesh.load_path(
                                    np.hstack((ray_origins, ray_origins + ray_directions / 2)).reshape(-1, 2, 3))
                                grasp_axis = transform[:3,1]
                                ray_axis =trimesh.load_path(
                                    np.hstack((transform[:3, 3].reshape(1, 3), transform[:3, 3].reshape(1, 3) + grasp_axis.reshape(1,3) / 5)).reshape(-1, 2, 3))
                                ray_axis_ =trimesh.load_path(
                                    np.hstack((transform[:3, 3].reshape(1, 3), transform[:3, 3].reshape(1, 3) - grasp_axis.reshape(1,3) / 5)).reshape(-1, 2, 3))
                                change_ray_color(ray_axis,[255,0,0])
                                change_ray_color(ray_axis_, [255, 0, 0])
                                scene_.add_geometry(ray_visualize, 'ray')
                                scene_.add_geometry(ray_axis, 'ray_axis')
                                scene_.add_geometry(ray_axis_, 'ray_axis_')
                                # scene_.add_geometry(gripper_copy, node_name='gripper')
                                scene_.add_geometry(pc_center, node_name='pc_center')
                                scene_.add_geometry(base_copy, node_name='base')
                                scene_.add_geometry(finger_left_copy, node_name='finger_left_copy')
                                scene_.add_geometry(finger_right_copy, node_name='finger_right_copy')
                                scene_.add_geometry(point_cloud, node_name='point_cloud')
                                if not is_collision and len(points_query) > 0:
                                    contact_points_cloud = trimesh.PointCloud(point_data[points_query],
                                                                              colors=[255, 0, 0])
                                    scene_.add_geometry(contact_points_cloud, node_name='contact_points_cloud')
                                show(scene_)
                                # scene_.delete_geometry('gripper')
                                scene_.delete_geometry('ray_axis')
                                scene_.delete_geometry('ray_axis_')
                                scene_.delete_geometry('pc_center')
                                scene_.delete_geometry('point_cloud')
                                scene_.delete_geometry('ray')
                                scene_.delete_geometry('base')
                                scene_.delete_geometry('finger_left_copy')
                                scene_.delete_geometry('finger_rifht_copy')
                                if not is_collision:
                                    scene_.delete_geometry('contact_points_cloud')
                                del scene_

                            # del gripper_copy
                            del finger_left_copy
                            del finger_right_copy
                            del base_copy
                            # print('test a single grasp angle')
                            # if not is_collision:
                            #     break
                    elif suction:

                        for grasp in grasps:
                            is_collision_flag = False
                            pose = RigidTransform(grasp.unrotated_full_axis, grasp.center, from_frame='gripper',
                                                  to_frame='obj')
                            transform = trimesh.transformations.concatenate_matrices(value, pose.matrix)
                            z_vec = -transform[:3, 2]

                            if z_vec[2] > -math.sin(math.pi / 4):
                                # print('test a single grasp angle')
                                is_collision_flag = True
                                # continue
                            suction_gripper_copy = deepcopy(suction_gripper)
                            suction_gripper_copy.apply_transform(transform)

                            if is_collision_flag:
                                is_collision = True
                            else:
                                is_collision = collision_manager.in_collision_single(suction_gripper_copy,
                                                                                     return_names=False)
                            if not is_collision:
                                contact_point = grasp.center
                                contact_point = trimesh.transformations.transform_points(contact_point[np.newaxis, ...],
                                                                                         value)
                                # contact_points_cloud = trimesh.PointCloud(contact_points, colors=[255, 0, 0])
                                points_query = point_tree.query_ball_point(contact_point, r=0.00175, p=2)[0]

                                if len(points_query) > 0:
                                    # the default setting let the points_query not appear in background or other object
                                    signed_distance = trimesh.proximity.signed_distance(mesh_dict[key_ori],
                                                                                        point_data[points_query])
                                    mask = abs(signed_distance) < 0.0005
                                    points_query = [points_query[i] for i in range(len(mask)) if mask[i]]
                                    for index in points_query:
                                        if index not in point_grasp_dict.keys() or grasp.score > \
                                                point_grasp_dict[index][1] or point_grasp_dict[index][0] == 0:
                                            point_grasp_dict[index] = (1, grasp.score, -1, -1, -1, 1)
                                            point_label_arr[index] = [1, grasp.score, -1, -1, -1, 1]
                            else:
                                contact_point = grasp.center
                                contact_point = trimesh.transformations.transform_points(contact_point[np.newaxis, ...],
                                                                                         value)
                                # contact_points_cloud = trimesh.PointCloud(contact_point, colors=[255, 0, 0])
                                # scene_tmp = trimesh.Scene([contact_points_cloud, point_cloud])
                                # scene_tmp.show()

                                points_query = point_tree.query_ball_point(contact_point, r=0.00175, p=2)[0]
                                # print(points_query)

                                if len(points_query) > 0:
                                    # the default setting let the points_query not appear in background or other object
                                    signed_distance = trimesh.proximity.signed_distance(mesh_dict[key_ori],
                                                                                        point_data[points_query])
                                    mask = abs(signed_distance) < 0.0005
                                    points_query = [points_query[i] for i in range(len(mask)) if mask[i]]
                                    for index in points_query:
                                        if index not in point_grasp_dict.keys() or (
                                                grasp.score > point_grasp_dict[index][1] and
                                                point_grasp_dict[index][0] == 0):
                                            point_grasp_dict[index] = (0, grasp.score, -1, -1, -1, 1)
                                            point_label_arr[index] = [0, grasp.score, -1, -1, -1, 1]

                            if debug_vis and not is_collision:
                                # if debug_vis:
                                # if debug_vis:
                                scene_ = deepcopy(scene)
                                # ray_origins = np.array([0., 0., 0.])
                                ray_origins = contact_point
                                ray_directions = z_vec
                                ray_visualize = trimesh.load_path(
                                    np.hstack((ray_origins, ray_origins + ray_directions / 2)).reshape(-1, 2, 3))
                                scene_.add_geometry(ray_visualize, 'ray')
                                # scene_.add_geometry(gripper_copy, node_name='gripper')
                                scene_.add_geometry(suction_gripper_copy, node_name='suction_gripper')
                                scene_.add_geometry(point_cloud, node_name='point_cloud')
                                if not is_collision and len(points_query) > 0:
                                    contact_points_cloud = trimesh.PointCloud(point_data[points_query],
                                                                              colors=[255, 0, 0])
                                    scene_.add_geometry(contact_points_cloud, node_name='contact_points_cloud')
                                show(scene_)
                                # scene_.delete_geometry('gripper')
                                scene_.delete_geometry('point_cloud')
                                scene_.delete_geometry('ray')
                                scene_.delete_geometry('suction_gripper')
                                if not is_collision and len(points_query) > 0:
                                    scene_.delete_geometry('contact_points_cloud')
                                del scene_

                            # del gripper_copy
                            del suction_gripper_copy
                            # print('test a single grasp angle')
                            # if not is_collision:
                            #     break
        # for point_index in obj_point_index:
        #     point_mask[point_index] = point_data[point_index] - value[:3, 3]
        #     print('point_direction :', point_mask[point_index])

    # scene_tmp.show()
    # scene.show()
    del collision_manager
    del scene
    print('test a single scene')
    # save_dict(point_grasp_dict, point_center_direction, k)
    save_arr(point_label_arr, k)


if __name__ == "__main__":
    debug_vis = True
    debug_pool = False
    use_pool = True

    if debug_vis:
        use_pool = False

    dst_path = '../grasp_suction_dataset/multi_obj_dataset'

    f = open(os.path.join(dst_path, 'metadata.json'), 'r')
    metadata = json.load(fp=f)
    f.close()

    suction_gripper_path = '/home/v-wewei/code/grasp_and_suction/grippers_model/yumi/suction_gripper.obj'
    base_path = '/home/v-wewei/code/grasp_and_suction/grippers_model/yumi/coarse/box_base.stl'
    finger_path = '/home/v-wewei/code/grasp_and_suction/grippers_model/yumi/coarse/new_box_finger.stl'  #origin finger is box_finger
    suction_gripper = trimesh.load_mesh(suction_gripper_path)
    base = trimesh.load_mesh(base_path)
    finger = trimesh.load_mesh(finger_path)
    T_gripper = trimesh.transformations.translation_matrix([0, 0, 0])
    R_gripper = trimesh.transformations.euler_matrix(np.pi / 2, np.pi / 2, 0, 'rxyz')
    T_base = trimesh.transformations.translation_matrix([0, 0, -0.161])
    R_base = trimesh.transformations.euler_matrix(np.pi / 2, np.pi / 2, 0, 'rxyz')
    T_finger = trimesh.transformations.translation_matrix([0, -0.0065, -0.048])

    suction_gripper.apply_transform(T_gripper)
    finger.apply_transform(T_finger)
    base.apply_transform(T_base)
    base.apply_transform(R_base)

    bin_mesh = trimesh.load_mesh('/home/v-wewei/code/grasp_and_suction/bin_model/bin.obj')
    bin_mesh.visual.face_colors = [60, 30, 78, 240]
    parallel_jaw_grasp = False
    suction = True
    assert not (parallel_jaw_grasp and suction)

    if parallel_jaw_grasp:
        grasp_path = '/home/v-wewei/code/grasp_and_suction/grasp_suction_dataset/grasp_{}/'.format(10000)
    elif suction:
        grasp_path = '/home/v-wewei/code/grasp_and_suction/grasp_suction_dataset/suction_{}/'.format(15000)

    if use_pool:
        pool = mp.Pool(processes=mp.cpu_count())
        #  pool = mp.Pool(processes=1)

    for i in range(0, 50):
        obj_id_file = np.load(os.path.join(dst_path, 'state_tensors/tensors/obj_ids_{:05d}.npz'.format(i)))
        obj_ids = obj_id_file['arr_0.npy'].reshape(-1, 15, 1)
        obj_pose_file = np.load(os.path.join(dst_path, 'state_tensors/tensors/obj_poses_{:05d}.npz'.format(i)))
        obj_poses = obj_pose_file['arr_0.npy'].reshape(-1, 15, 7)

        for j in range(2):
            depth_tensor = \
                np.load(os.path.join(dst_path, 'image_tensors/tensors/depth_im_{:05d}.npz'.format(i * 2 + j)))[
                    'arr_0.npy']
            camera_intrs_tensor = \
                np.load(os.path.join(dst_path, 'image_tensors/tensors/camera_intrs_{:05d}.npz'.format(i * 2 + j)))[
                    'arr_0.npy']
            camera_pose_tensor = \
                np.load(os.path.join(dst_path, 'image_tensors/tensors/camera_pose_{:05d}.npz'.format(i * 2 + j)))[
                    'arr_0.npy']
            modal_mask_tensor = np.load(os.path.join(dst_path, 'image_tensors/tensors/modal_segmasks_{:05d}.npz'.format(i * 2 + j)))['arr_0.npy']
            k = 200 * i + j * 100
            for image_data, camera_intr_, camera_pose, modal_mask in zip(depth_tensor, camera_intrs_tensor, camera_pose_tensor,
                                                                         modal_mask_tensor):
                if k % 2 != 0:
                    k += 1
                    continue
                if use_pool:
                    if debug_pool:
                        pool.apply(generate_label, args=(k, suction_gripper, base, finger, bin_mesh,
                                                         obj_ids, obj_poses, image_data, camera_intr_, camera_pose,
                                                         modal_mask))
                    else:
                        last = pool.apply_async(generate_label, args=(k, suction_gripper, base, finger, bin_mesh, obj_ids,
                                                               obj_poses, image_data, camera_intr_, camera_pose,
                                                               modal_mask))
                        if len(pool._cache) > 8e2:
                            print('waiting for cache to clear')
                            last.wait()

                else:
                    generate_label(k, suction_gripper, base, finger, bin_mesh, obj_ids, obj_poses, image_data,
                                   camera_intr_, camera_pose, modal_mask)
                k += 1
    if use_pool:
        pool.close()
        pool.join()
